# %%
import tensorflow as tf
from tensorflow import keras
import numpy as np
import cv2
from PIL import Image
from queue import Queue
from concurrent import futures
import time

# %%
p2v = keras.models.load_model('pic2vec.h5')
p2v.summary()


# %%
def gen_ds(fpath):
  vid = cv2.VideoCapture(fpath)
  ok, frame = vid.read()
  ds = []
  counter = 1
  while ok:
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    img = Image.fromarray(frame)
    w, h = img.size
    r = max(w, h)
    img = img.crop((0, 0, r, r))
    # img = img.resize((32, 32))
    counter += 1
    if counter % 100 == 0:
      print(counter)
    ds.append(np.asarray(img))
    ok, frame = vid.read()
  return ds


# %%
ds = gen_ds('test/test-32.mkv')
# %%
ds = np.asarray(ds)
print(ds.shape)
# %%
res = p2v.predict(ds)
# %%
# print(np.linalg.norm(res[150] - res[200]))

arr_len = len(res)
groups = []
used = []
for x in range(arr_len):
  if x in used:
    continue
  group = []
  group.append(x)
  for y in range(x + 1, arr_len):
    if y in used:
      continue
    dist = np.linalg.norm(res[x] - res[y])
    if dist < 30:
      used.append(y)
      group.append(y)
  if len(group) > 1:
    groups.append(np.asarray(group))
  # if len(groups) > 10:
    # break

# %%
groups = np.asarray(groups)
print(groups.shape)
# %%
print(len(groups[-1]))